import os
import re
import json
import datetime
from openai import OpenAI

# ===== 路径配置 =====

ROOT_DIR = r"E:\沙粒云\自媒体\2025视频制作\20250610ai高考\流程\anwser\data"
INPUT_FULL_TEXT = os.path.join(ROOT_DIR, "output_full_exam.txt")
ANWSER_DIR = os.path.join(ROOT_DIR, "anwser")
# 虽然不再输出 JSON 格式单元，这个常量可以保留或忽略
OUTPUT_UNIT_JSON = os.path.join(ANWSER_DIR, "unit_split.json")
OUTPUT_ANSWERS = os.path.join(ANWSER_DIR, "answer_list.json")
UNITS_FOLDER = os.path.join(ANWSER_DIR, "units")
UNIT_ANSWERS_FOLDER = os.path.join(ANWSER_DIR, "unit_answers")
LOG_FILE = os.path.join(ANWSER_DIR, "b_pipeline_log.txt")

# ===== Qwen-3 接口配置 =====

API_KEY      = "sk-c48712427ba84b59b9c89fc2883e48b6"
BASE_URL     = "https://dashscope.aliyuncs.com/compatible-mode/v1"
QWEN_3_MODEL = "qwen3-235b-a22b"

# 初始化 Qwen 客户端
try:
    qwen_client = OpenAI(api_key=API_KEY, base_url=BASE_URL)
except Exception as e:
    print(f"[ERROR] OpenAI 客户端初始化失败: {e}")
    raise


# ===== DeepSeek 接口配置 (用于答题) =====

DEEPSEEK_API_KEY = "sk-55bcb4efc5c4402d86c30c7eb8a3bdf5" # <--- 请务必替换为您自己的 DeepSeek Key
DEEPSEEK_BASE_URL = "https://api.deepseek.com"
DEEPSEEK_MODEL = "deepseek-coder" # 或者使用 "deepseek-reasoner"
# 初始化 DeepSeek 客户端
try:
    deepseek_client = OpenAI(api_key=DEEPSEEK_API_KEY, base_url=DEEPSEEK_BASE_URL)
except Exception as e:
    print(f"[ERROR] DeepSeek 客户端初始化失败: {e}")
    raise


def log(msg: str):
    """
    在控制台打印并写入日志文件 b_pipeline_log.txt
    """
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    line = f"[{timestamp}] {msg}"
    print(line)
    os.makedirs(os.path.dirname(LOG_FILE), exist_ok=True)
    with open(LOG_FILE, "a", encoding="utf-8") as f:
        f.write(line + "\n")

def load_full_text(path: str) -> str:
    """
    读取完整 OCR 文本（output_full_exam.txt）
    """
    if not os.path.isfile(path):
        log(f"❌ 文件不存在: {path}")
        raise FileNotFoundError(f"{path} 不存在")
    with open(path, "r", encoding="utf-8") as f:
        content = f.read()
    log(f"成功读取 OCR 文本，字符长度: {len(content)}")
    return content

def split_into_units_text(full_text: str) -> str:
    """
    不要求模型输出 JSON，而是让模型按以下格式输出纯文本：
      ==== Unit1 ====
      <单元1标题>
      <阅读材料原文>
      【第<index>题】
      <题干文本>
      （题与题之间留一个空行）
      ==== Unit2 ====
      <单元2标题>
      ……
    """
    log(f"进入 split_into_units_text，收到 full_text 长度: {len(full_text)} 字符")

    prompt = (
        "请将下面的 OCR 完整文本切分成若干“单元”，每个单元严格按照如下格式输出：\n"
        "==== UnitN ====\n"
        "<单元N的标题>\n"
        "<该单元阅读材料原文（跨页内容一并写出）>\n"
        "每道题以“【第<index>题】”开头，后跟题干原文，题与题之间留一个空行。\n"
        "不要输出多余说明，保持和原文一样只是分开，当单元跨了多页，可以删去单元中间的（第x页/共x页【第x页】）除此之外不要比原文增加或者减少任何一个字。只输出上述格式的纯文本。\n\n"
        f"{full_text}"
    )

    log("调用 Qwen-3 切割单元（split_into_units_text），开始发送请求")
    try:
        response = qwen_client.chat.completions.create(
            model=QWEN_3_MODEL,
            messages=[
                {"role": "system", "content": "你是一个文本切割助手，输出时严格按照约定格式。"},
                {"role": "user", "content": prompt}
            ],
            stream=True,
            stream_options={"include_usage": False}
        )
    except Exception as e:
        log(f"❌ split_into_units_text 请求异常: {e}")
        raise

    accumulated = ""
    chunk_count = 0
    log("流式调用已发出，开始接收流式返回")
    for chunk in response:
        chunk_count += 1
        try:
            chunk_data = json.loads(chunk.model_dump_json())
            delta = chunk_data["choices"][0].get("delta", {})
            piece = delta.get("content", "")
            if piece:
                accumulated += piece
                if chunk_count == 1:
                    preview = piece[:50].replace("\n", " ")
                    log(f"[split_text] 第1个 chunk 内容预览: {preview} …")
                print(piece, end="", flush=True)
        except Exception as e:
            log(f"[split_text] 解析第 {chunk_count} 个 chunk 时出错: {e}")
            continue

    if chunk_count == 0:
        log("❗ 警告：split_into_units_text 未收到任何 chunk")

    log(f"流式接收结束，共收到 {chunk_count} 个 chunk，累计字符数: {len(accumulated.strip())}")
    return accumulated.strip()

def save_units_from_text(raw: str):
    """
    将 split_into_units_text 返回的整段带分隔符的文本按 “==== UnitN ====”
    拆分，并将每个单元的内容（标题+阅读材料+题干列表）分别写入独立的 .txt 文件。
    """
    os.makedirs(UNITS_FOLDER, exist_ok=True)
    log("开始将模型返回的原始文本按单元拆分并保存")

    # 用正则匹配所有 “==== UnitN ====”
    pattern = re.compile(r"^====\s*(Unit\d+)\s*====\s*$", re.MULTILINE)
    matches = list(pattern.finditer(raw))

    if not matches:
        log("❌ 未找到任何单元标志 “==== UnitN ====”，无法拆分")
        return

    for idx, match in enumerate(matches):
        unit_name = match.group(1)  # e.g., "Unit1"
        start_pos = match.end()
        if idx + 1 < len(matches):
            end_pos = matches[idx + 1].start()
        else:
            end_pos = len(raw)
        body = raw[start_pos:end_pos].strip()

        filename = os.path.join(UNITS_FOLDER, f"{unit_name}.txt")
        try:
            with open(filename, "w", encoding="utf-8") as f:
                f.write(body)
            log(f"  已保存 {unit_name}.txt → {filename}")
        except Exception as e:
            log(f"  ❌ 保存 {unit_name}.txt 失败: {e}")

def get_unit_filenames(folder: str) -> list:
    """
    返回 folder 目录下所有 .txt 文件的绝对路径列表（按名称排序）
    """
    if not os.path.isdir(folder):
        log(f"❗ 单元目录不存在: {folder}")
        return []
    files = [f for f in os.listdir(folder) if f.lower().endswith(".txt")]
    files.sort()
    return [os.path.join(folder, f) for f in files]

def ask_model_for_answers(unit_text: str, unit_name: str) -> str:
    """
    将某个单元的完整文本发送给 Qwen-3，让它输出每道题的答案（每行一个“第N题: 答案”格式）。
    对于子题标签，要保留完整编号，例如 13（1）、13（2）。
    返回模型原始回复（文本）。
    """
    prompt = (
        f"你是一个严格遵守格式的答题机器人。请根据下方的题目文本，严格模仿示例格式，生成答案列表。\n\n"
        "### 输出格式示例 ###\n"
        "第1题: A\n"
        "第2题: B\n"
        "第13(1)题: 这是第一小题的答案。\n"
        "第13(2)题: 这是第二小题的答案。\n"
        "第14题: 汉皇重色思倾国\n"
        "### 输出格式示例结束 ###\n\n"
        f"### “{unit_name}”单元的题目文本 ###\n"
        f"{unit_text}\n"
        "### 题目文本结束 ###\n\n"
        "请开始作答，你的回答中除了答案列表外，不要包含任何其他内容。"
    )
    log(f"调用 DeepSeek 解答 {unit_name}")
    try:
        # 使用 deepseek_client 和 DEEPSEEK_MODEL
        response = deepseek_client.chat.completions.create(
            model=DEEPSEEK_MODEL,
            messages=[
                # System message 强调了遵守格式的重要性，并保留了作文要求
                {"role": "system", "content": "你是一个高考解题专家，请严格按照用户要求的格式，分析以下试卷题目并给出答案。当题目中出现“写作”“写一篇文章”，请按照满分标准写一篇大于1200字的文章。"},
                {"role": "user", "content": prompt}
            ],
            stream=True
        )
    except Exception as e:
        log(f"  ❌ {unit_name} 解题请求失败: {e}")
        return f"模型请求失败：{e}"

    accumulated = ""
    chunk_count = 0
    log(f"  [answer-{unit_name}] DeepSeek 流式调用已发出，开始接收返回...")
    for chunk in response:
        chunk_count += 1
        try:
            # DeepSeek 的流式数据结构更直接
            piece = chunk.choices[0].delta.content

            if piece:
                accumulated += piece
                if chunk_count == 1:
                    preview = piece[:50].replace("\n", " ")
                    log(f"  [answer-{unit_name}] 第1个 chunk 内容预览: {preview} …")
                print(piece, end="", flush=True)

        except Exception as e:
            log(f"  [answer-{unit_name}] 解析第 {chunk_count} 个 chunk 时出错: {e}")
            continue

    if chunk_count == 0:
        log(f"  ❗ 警告：{unit_name} 未收到任何解题 chunk")

    # 换行以确保日志格式清晰
    print() 

    ans = accumulated.strip()
    log(f"  {unit_name} 返回答案，共 {len(ans.splitlines())} 行")
    return ans

def main():
    # 准备目录
    os.makedirs(ANWSER_DIR, exist_ok=True)
    os.makedirs(UNITS_FOLDER, exist_ok=True)
    os.makedirs(UNIT_ANSWERS_FOLDER, exist_ok=True)

    # 初始化日志文件
    with open(LOG_FILE, "w", encoding="utf-8") as f:
        f.write(f"[{datetime.datetime.now()}] 全流程脚本启动\n")

    # 1. 读取 OCR 文本
    log("读取完整试卷文本")
    try:
        full_text = load_full_text(INPUT_FULL_TEXT)
    except Exception as e:
        log(f"❌ 读取 OCR 文本失败：{e}")
        return

    # 2. 调用模型生成带分隔符的单元文本
    log("调用模型生成带分隔符的单元文本")
    try:
        raw_units_text = split_into_units_text(full_text)
    except Exception as e:
        log(f"❌ split_into_units_text 执行失败: {e}")
        return

    # 3. 按分隔符拆分并保存为 .txt 文件
    save_units_from_text(raw_units_text)

    # 4. 逐单元读取 .txt 并让模型解答，保存答案
    unit_files = get_unit_filenames(UNITS_FOLDER)
    if not unit_files:
        log("未找到任何单元 txt，脚本退出。")
        return

    all_answers = []
    for unit_path in unit_files:
        unit_name = os.path.splitext(os.path.basename(unit_path))[0]
        log(f"开始处理 {unit_name}")
        try:
            with open(unit_path, "r", encoding="utf-8") as f:
                unit_text = f.read().strip()
        except Exception as e:
            log(f"❌ 读取单元文件 {unit_path} 失败: {e}")
            continue

        ans_text = ask_model_for_answers(unit_text, unit_name)

        answer_file = os.path.join(UNIT_ANSWERS_FOLDER, f"{unit_name}_answer.txt")
        try:
            with open(answer_file, "w", encoding="utf-8") as f:
                f.write(ans_text)
            log(f"  已保存 {unit_name}_answer.txt → {answer_file}")
        except Exception as e:
            log(f"  ❌ 保存 {unit_name}_answer.txt 失败: {e}")

        ans_lines = [line.strip() for line in ans_text.splitlines() if line.strip()]
        all_answers.append({"unit": unit_name, "answers": ans_lines})

    # 5. 汇总所有单元答案到 answer_list.json
    try:
        with open(OUTPUT_ANSWERS, "w", encoding="utf-8") as f:
            json.dump(all_answers, f, ensure_ascii=False, indent=2)
        log(f"所有单元答案汇总已保存到：{OUTPUT_ANSWERS}")
    except Exception as e:
        log(f"❌ 保存 answer_list.json 失败: {e}")
        return

    log("全流程脚本执行完毕")

if __name__ == "__main__":
    main()
